# -*- coding: utf-8 -*-
"""
Created on Wed Mar 13 21:04:48 2024

@author: fkxxgis

https://blog.csdn.net/zhebushibiaoshifu/article/details/138233711
"""

import os
import fnmatch
import numpy as np
import matplotlib.pyplot as plt
from osgeo import gdal
from matplotlib.colors import LinearSegmentedColormap

file_path_tif = r"D:\GPPvsWater\WaterYield1981-2022\tiffs\wy_0p5deg"
file_path_plot = r"D:\GPPvsWater\WaterYield1981-2022"
file_extension = "*.tif"
band_idx = 1;
row = 6
col = 8

file_all = os.listdir(file_path_tif)
file_all_tif = [file_tif for file_tif in file_all if fnmatch.fnmatch(file_tif, file_extension)]
file_all_tif.sort()

if band_idx == 1:
    cmap_colors = [(0, "white"), (0.1, "powderblue"), (0.2, "lightblue"), (0.3, "lightskyblue"), (0.4, "skyblue"),
                   (0.5, "deepskyblue"), (0.6, "cornflowerblue"), (0.7, "royalblue"), (0.8, "blue"),
                   (0.9, "mediumblue"), (1, "darkblue")]
    cmap = LinearSegmentedColormap.from_list("", cmap_colors)
    file_pic_name = "blue_ori.png"
elif band_idx == 2:
    cmap_colors = [(0, "white"), (0.1, "palegreen"), (0.2, "lightgreen"), (0.3, "greenyellow"), (0.4, "lawngreen"),
                   (0.5, "springgreen"), (0.6, "lime"), (0.7, "limegreen"), (0.8, "forestgreen"),
                   (0.9, "green"), (1, "darkgreen")]
    cmap = LinearSegmentedColormap.from_list("", cmap_colors)
    file_pic_name = "green_ori.png"
elif band_idx == 3:
    cmap_colors = [(0, "white"), (0.1, "darksalmon"), (0.2, "salmon"), (0.3, "tomato"), (0.4, "orangered"),
                   (0.5, "red"), (0.6, "indianred"), (0.7, "firebrick"), (0.8, "brown"),
                   (0.9, "maroon"), (1, "darkred")]
    cmap = LinearSegmentedColormap.from_list("", cmap_colors)
    file_pic_name = "red_ori.png"
elif band_idx == 4:
    cmap_colors = [(0, "white"), (0.1, "ivory"), (0.2, "lemonchiffon"), (0.3, "papayawhip"), (0.4, "navajowhite"),
                   (0.5, "sandybrown"), (0.6, "orange"), (0.7, "darkorange"), (0.8, "peru"),
                   (0.9, "chocolate"), (1, "saddlebrown")]
    # cmap = LinearSegmentedColormap.from_list("", cmap_colors)
    cmap = "YlOrRd"
    file_pic_name = "NIR_ori_2.png"
else:
    cmap = "Greens"
    file_pic_name = "NDVI_MODIS.png"

idx = 1
plt.rc("font", family="Times New Roman")
fig = plt.figure(figsize=(45, 20))

for file_tif in file_all_tif:
    dataset = gdal.Open(os.path.join(file_path_tif, file_tif))
    band = dataset.GetRasterBand(band_idx)
    array_band = band.ReadAsArray()
    # nodata = band.GetNoDataValue()
    array_band = np.nan_to_num(array_band, nan=0)
    if np.all(array_band == 0):
        img = plt.imshow(np.zeros((1, 1)), cmap=cmap)
        plt.axis("off")
        plt.title(file_tif, y=-0.2)
        idx += 1
        continue
    array_band[(array_band > 10000) | (array_band < 0)] = 0
    array_band = array_band / 1.0
    dataset = None

    ax = plt.subplot(row, col, idx)
    img = plt.imshow(array_band, cmap=cmap)
    plt.axis("off")
    plt.title(file_tif, y=-0.2)
    print(idx, "finished!")
    idx += 1

cax = plt.axes([0.66, 0.15, 0.24, 0.015])
colorbar = plt.colorbar(cax=cax, orientation="horizontal")
colorbar.outline.set_visible(False)
plt.savefig(file_path_plot + "/" + file_pic_name, dpi=300, bbox_inches="tight", pad_inches=0)
plt.show()
